package org.example;

import java.util.Random;

public class Main {
    public static void main(String[] args) {
        System.out.println("Hello world!");
    }
}
class Solution {
    public int[][] kClosest(int[][] points, int k) {
        quickSelect(points, 0, points.length - 1, k);

        //返回前k个元素
        int[][] res = new int[k][2];
        for (int i = 0; i < k; i++) {
            res[i] = points[i];
        }
        return res;
    }

    private void quickSelect(int[][] points, int l, int r, int k) {
        if(l >= r) return;

        //随机选择基准值
        int[] key = points[new Random().nextInt(r - l + 1) + l];
        int keyDistance = distance(key);
        //三路划分
        int left = l - 1, right = r + 1, i = l;
        while(i < right){
            if(distance(points[i]) < keyDistance){
                Swap(points, ++left, i++);
            }else if(distance(points[i]) == keyDistance){
                i++;
            }else {
                Swap(points, --right, i);
            }
        }
        //[l, left] [left + 1, right - 1] [right, r]
        int a = left - l + 1, b = right - left - 1;
        if(a >= k)  quickSelect(points, l, left, k);
        else if(a + b >= k) return;
        else   quickSelect(points, right, r, k - a - b);
    }

    private void Swap(int[][] points, int i, int j) {
        int[] t = points[i];
        points[i] = points[j];
        points[j] = t;
    }

    //计算距离原点的距离,避免开根号
    private int distance(int[] key) {
        return key[0] * key[0] + key[1] * key[1];
    }

}